#ifndef YCSB_C_DHMS_DB_H_
#define YCSB_C_DHMS_DB_H_

#include "client/client.h"
#include "rtx/txn.h"
#include "rtx/cort_pool.h"
#include "core/db.h"
#include "core/properties.h"
#include "core/core_workload.h"
#include "core/timer.h"
#include <atomic>
#include <iostream>
#include <vector>

using namespace dhms;

namespace ycsbc {

class DhmsDB : public DB {
    class AllocTask : public rtx::Task {
    public: 
        AllocTask(PoolClient* client_hdlr, DcObject** dest, 
                  size_t obj_size, std::atomic<int>* cnter) {
            assert(dest);
            txn_ = new rtx::Txn(client_hdlr);
            dest_ = dest;
            obj_size_ = obj_size;
            cnter_ = cnter;
        }

        void operator()() override {
            // printf("Alloc Task Begin!\n");
            txn_->begin();
            rtx::Txn::TXN_EXEC_STATUS rc = txn_->Alloc(obj_size_, dest_);
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);
            txn_->commit();
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);

            cnter_->fetch_add(1);
            // printf("Alloc Task Done!\n");
        }

        rtx::Txn* txn_;
        DcObject** dest_;
        size_t obj_size_;
        std::atomic<int>* cnter_;
    };

    class InsertTask : public rtx::Task {
    public:
        InsertTask(PoolClient* client_hdlr, DcObject* dest, 
                  size_t obj_size, std::atomic<int>* cnter,
                  const std::string& key) {
            assert(dest);
            txn_ = new rtx::Txn(client_hdlr);
            dest_ = dest;
            obj_size_ = obj_size;
            cnter_ = cnter;
            cbuf_ = new char[key.size()];
            len_ = key.size();
            memcpy(cbuf_, key.c_str(), key.size());
        }

        void operator()() override {
            txn_->begin();
            rtx::Txn::ObjOffsetBuf buf;
            d_iov_set(&(buf.buf), cbuf_, len_);
            buf.offset = 0;
            rtx::Txn::TXN_EXEC_STATUS rc = txn_->WriteOnce(dest_, &buf);
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);
            rc = txn_->commit();
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);

            cnter_->fetch_add(1);
        }        

        rtx::Txn* txn_;
        DcObject* dest_;
        size_t obj_size_;
        std::atomic<int>* cnter_;
        char* cbuf_;
        size_t len_;
    };

    class ReadTask : public rtx::Task {
    public:
        ReadTask(PoolClient* client_hdlr, DcObject* dest, 
                  size_t obj_size, std::atomic<int>* cnter) {
            assert(dest);
            txn_ = new rtx::Txn(client_hdlr);
            dest_ = dest;
            obj_size_ = obj_size;
            cnter_ = cnter;
        }

        void operator()() override {
            txn_->begin();
            char* read_buf = (char*)global_rrpc_ring_buffer->Alloc(obj_size_);
            rtx::Txn::ObjOffsetBuf read_obj_buf;
            d_iov_set(&(read_obj_buf.buf), read_buf, obj_size_);
            read_obj_buf.offset = 0;
            rtx::Txn::TXN_EXEC_STATUS rc = txn_->ReadOnce(dest_, &read_obj_buf);
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);
            rc = txn_->commit();
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);

            // std::string str(read_buf, obj_size_);
            // printf("read: %s\n", str.c_str());
            global_rrpc_ring_buffer->Free(read_buf);
            cnter_->fetch_add(1);
        }

        rtx::Txn* txn_;
        DcObject* dest_;
        size_t obj_size_;
        std::atomic<int>* cnter_;
    };

    class RMWTask : public rtx::Task {
    public:
        RMWTask(PoolClient* client_hdlr, DcObject* dest, 
                size_t obj_size, std::atomic<int>* cnter,
                const std::string& key) {
            assert(dest);
            txn_ = new rtx::Txn(client_hdlr);
            dest_ = dest;
            obj_size_ = obj_size;
            cnter_ = cnter;
            cbuf_ = new char[key.size()];
            len_ = key.size();
            memcpy(cbuf_, key.c_str(), key.size());
        }

        void operator()() override {
            txn_->begin();
            // Read
            char* read_buf = (char*)global_rrpc_ring_buffer->Alloc(obj_size_);
            rtx::Txn::ObjOffsetBuf read_obj_buf;
            d_iov_set(&(read_obj_buf.buf), read_buf, obj_size_);
            read_obj_buf.offset = 0;
            rtx::Txn::TXN_EXEC_STATUS rc = txn_->ReadOnce(dest_, &read_obj_buf);
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);
            // Modify
            rtx::Txn::ObjOffsetBuf buf;
            d_iov_set(&(buf.buf), cbuf_, len_);
            buf.offset = 0;
            rc = txn_->WriteOnce(dest_, &buf);
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);
            // Write
            rc = txn_->commit();
            // assert(rc == rtx::Txn::TXN_EXEC_SUCCESS);

            global_rrpc_ring_buffer->Free(read_buf);
            cnter_->fetch_add(1);
        }

        rtx::Txn* txn_;
        DcObject* dest_;
        size_t obj_size_;
        std::atomic<int>* cnter_;
        char* cbuf_;
        size_t len_;
    };
public:
    DhmsDB(size_t record_cnt, int thread_num, utils::Properties& props) : 
        obj_cnt_(record_cnt), thread_num_(thread_num) {
        dev_id_ = (uint8_t)stoi(props.GetProperty("dev_id", "0"));
        port_id_ = (uint8_t)stoi(props.GetProperty("port_id", "2"));
        rrpc_thread_num_ = (uint8_t)stoi(props.GetProperty("rrpc_thread_num", "4"));
        client_ip_ = props.GetProperty("client_ip", "192.168.1.22");
        client_port_ = stoll(props.GetProperty("client_port", "10070"));
        client_start_port_ = stoll(props.GetProperty("client_start_port", "10070"));
        master_ip_ = props.GetProperty("master_ip", "192.168.1.22");
        master_port_ = stoll(props.GetProperty("master_port", "10086"));
        field_len_ = stoi(props.GetProperty(CoreWorkload::FIELD_LENGTH_PROPERTY, 
                                            CoreWorkload::FIELD_LENGTH_DEFAULT));

        printf("============[DhmsDB Config]============\n"\
               "dev_id: %u\n"\
               "port_id: %u\n"\
               "rrpc_thread_num: %u\n"\
               "client_ip: %s\n"\
               "client_port: %lu\n"\
               "client_start_port: %lu\n"\
               "master_ip: %s\n"\
               "master_port: %lu\n"\
               "thread_num: %d\n"\
               "field_len: %d\n\n", 
               dev_id_, port_id_, rrpc_thread_num_, 
               client_ip_.c_str(), client_port_, client_start_port_,
               master_ip_.c_str(), master_port_, thread_num_,
               field_len_);
    }

    void Init() {
        printf("************** DB->Init ***************\n");
        dhms_server_id_t sid = 0;
        // printf("ip: %s port: %lu sid: %lu\n", client_ip_.c_str(), client_port_, sid);
        rdma_ctrl_ = new rpc::RDMACtrl(client_ip_, client_port_, sid, true);
        rpc::RNicHandler* rnic = rdma_ctrl_->open_dev(dev_id_, port_id_);

        rrpc_options_t rrpc_options;
        rrpc_options.need_thread_pool = false;
        global_rrpc_mgr = new rpc::RRPCManager(rrpc_thread_num_, rdma_ctrl_, rnic, master_ip_,
                                                master_port_, 0, true, rrpc_options);
        
        dhms_server_id_t master_leader_sid = global_rrpc_mgr->get_master_id();
        ClientAgent* client_agent = new ClientAgent(master_leader_sid, client_ip_.c_str(), client_port_);
        int ret_error = 0;
        dhms_pool_uuid_t pool_uuid = 1;
        std::string pool_name = "DFS";
        size_t size = 100UL * 1024 * 1024 * 1024;

        // PoolClient* pool_client = nullptr;
        ret_error = client_agent->PoolCreate(pool_name, size, &pool_client_);
        pool_uuid = pool_client_->GetPoolUUID();

        pool_ = new rtx::CortPool(thread_num_);

        all_objs_.resize(obj_cnt_);
        std::atomic<int> alloc_cnter{0};
        utils::Timer<double> timer;
        timer.Start();
        for (int i = 0; i < obj_cnt_; ++i) {
            pool_->add_task(new AllocTask(pool_client_, &(all_objs_[i]), field_len_, &alloc_cnter));
        }

        while (alloc_cnter.load() < obj_cnt_) {
            // printf("alloc_cnter: %d\n", alloc_cnter.load());
            cpu_relax();
        }
        double duration = timer.End();
        printf("+++++++++++++ Dhms DB init OK! ++++++++++++\n");
#ifndef NDEBUG
        std::set<dhms_oid_t> oid_set;
        for (int i = 0; i < obj_cnt_; ++i) {
            // printf("Oid[%d]: %lu\n", i, all_objs_[i]->GetOID());
            dhms_oid_t cur_oid = all_objs_[i]->GetOID();
            if (oid_set.find(cur_oid) != oid_set.end()) {
                printf("!!!!!err!!!!! find overlaped oid: %lu", cur_oid);
            }
            oid_set.insert(cur_oid);
        }
        printf("+++++++++++++ Dhms DB Check PASS! ++++++++++++\n");
#endif
        std::cerr << "# Alloc OID throughput " << (obj_cnt_ / duration / 1000) << " (KTPS)" << std::endl;
    }

    int Read(const std::string &table, const std::string &key,
           const std::vector<std::string> *fields,
           std::vector<KVPair> &result) {
        return 0;
    }

    int ReadAsync(const std::string &table, const std::string &key,
                  const std::vector<std::string> *fields,
                  std::vector<KVPair> &result, std::atomic<int>* cnter) 
    {
        DcObject* obj = all_objs_[stoi(table)];
        pool_->add_task(new ReadTask(pool_client_, obj, field_len_, 
                                    cnter));
        return 0;
    }

    int Scan(const std::string &table, const std::string &key,
           int len, const std::vector<std::string> *fields,
           std::vector<std::vector<KVPair>> &result) {
        return 0;
    }

    int Update(const std::string &table, const std::string &key,
             std::vector<KVPair> &values) {
        return 0;
    }

    int UpdateAsync(const std::string &table, const std::string &key,
                     std::vector<KVPair> &values, std::atomic<int>* cnter) {
        DcObject* obj = all_objs_[stoi(table)];
        pool_->add_task(new InsertTask(pool_client_, obj, field_len_, 
                                        cnter, key));
        return 0;
    }

    int RMWAsync(const std::string &table, const std::string &key,
                 std::vector<KVPair> &values, std::atomic<int>* cnter) {
        DcObject* obj = all_objs_[stoi(table)];
        pool_->add_task(new RMWTask(pool_client_, obj, field_len_, 
                                    cnter, key));
        return 0;
    }

    int Insert(const std::string &table, const std::string &key,
             std::vector<KVPair> &values) {
        return 0;
    }

    int InsertAsync(const std::string &table, const std::string &key,
                     std::vector<KVPair> &values, std::atomic<int>* cnter) {
        DcObject* obj = all_objs_[stoi(table)];
        pool_->add_task(new InsertTask(pool_client_, obj, field_len_, 
                                        cnter, key));
        return 0;
    }

    int Delete(const std::string &table, const std::string &key) {
        return 0;
    }

private:
    uint8_t dev_id_;
    uint8_t port_id_;
    // const std::string rrpc_pm_file = "/dev/dax0.1";
    // uint64_t rrpc_pm_offset = 5UL * 1024 * 1024 * 1024;
    // uint64_t rrpc_pm_len = 5UL * 1024 * 1024 * 1024;
    uint8_t rrpc_thread_num_;
    std::string client_ip_;
    uint64_t client_port_;
    uint64_t client_start_port_;
    std::string master_ip_;
    uint64_t master_port_;
    int field_len_;

    rpc::RDMACtrl *rdma_ctrl_ = nullptr;
    rtx::CortPool* pool_ = nullptr;
    PoolClient* pool_client_ = nullptr;

    size_t obj_cnt_;
    int thread_num_;
    std::vector<DcObject*> all_objs_;
};

}

#endif